import joblib
import pandas as pd
import torch
from torch import nn
import torch.nn.functional as F

from predicte.train import scaler


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv1d(in_channels=9, out_channels=64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv1d(in_channels=64, out_channels=128, kernel_size=3, padding=1)
        self.maxpool = nn.MaxPool1d(2)
        self.lstm1 = nn.LSTM(input_size=64, hidden_size=128, batch_first=True)
        self.dropout1 = nn.Dropout(0.2)
        self.bidirectional = nn.LSTM(128, 128, bidirectional=True, batch_first=True)
        self.dropout2 = nn.Dropout(0.2)
        self.dense1 = nn.Linear(128 * 2, 64)
        self.dense2 = nn.Linear(64, 8)
        self.dense3 = nn.Linear(8, 1)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = self.maxpool(x.permute(0, 2, 1))
        x, _ = self.lstm1(x)
        x = self.dropout1(x)
        x, _ = self.bidirectional(x)
        x = x[:, -1, :]
        x = self.dropout2(x)
        x = F.sigmoid(self.dense1(x))
        x = self.dense2(x)
        x = self.dense3(x)
        return x


Torchmodel = Net()  # 确保Net类定义与训练时一致
Torchmodel.load_state_dict(torch.load('model_state_dict.pth'))
Torchmodel.eval()  # 设置为评估模式


# 如果你保存了整个模型，可以直接加载
# Torchmodel = torch.load('model.pth')
# Torchmodel.eval()

# 假设你有一个函数来准备输入数据
def prepare_input_data(date, df, scaler, lookback):
    idx = df[df['DATE'] == date].index[0]
    feature_columns = df.columns
    print(feature_columns)
    if 'DATE' in feature_columns:
        feature_columns = feature_columns.drop(['DATE', 'maxTemp', 'year'])
        print(feature_columns)
    feature_data = df.loc[idx - lookback + 1:idx, feature_columns].values
    print(feature_data.shape)
    feature_data = scaler.transform(feature_data.astype('float32'))  # 归一化

    # 准备输入张量
    input_tensor = torch.tensor(feature_data).unsqueeze(0)  # 增加batch维度
    return input_tensor


def predict_future_dates(model, df, scaler, start_date, num_days, lookback=3):
    model.eval()  # 设置为评估模式
    predictions = []
    start_prediction_date = pd.to_datetime(start_date)

    for i in range(num_days):
        prediction_date = start_prediction_date + pd.Timedelta(days=i)
        print(prediction_date)
        print(df)
        if prediction_date not in df['DATE'].values:
            pass
        try:
            input_tensor = prepare_input_data(prediction_date, df, scaler, lookback)
            with torch.no_grad():
                prediction = model(input_tensor.permute(0, 2, 1))
                prediction = prediction.item()
            predictions.append((prediction_date, prediction))
        except Exception as e:
            print(f"Error processing date {prediction_date}: {e}")
            predictions.append((prediction_date, None))

    return predictions


# 使用示例
def test():
    max = []
    # 加载模型和 scaler
    Torchmodel = Net()
    Torchmodel.load_state_dict(torch.load('model_state_dict.pth', map_location=torch.device('cpu')))
    # scaler = joblib.load('scaler.pkl')

    # 加载数据
    df = pd.read_csv("new_data.csv")
    df['DATE'] = pd.to_datetime(df['DATE'])
    print(df)

    # 进行预测
    start_date = '2024-01-05'
    num_days = 7
    predictions = predict_future_dates(Torchmodel, df, scaler, start_date, num_days)

    # 打印预测结果
    for date, pred in predictions:
        if pred is not None:
            max.append(pred)
            print(f"Prediction for {date}: {pred}")
        else:
            print(f"No prediction available for {date}.")
    return max
